{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Separate Total Population dataset by States"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cudf\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>easting</th>\n",
       "      <th>northing</th>\n",
       "      <th>race</th>\n",
       "      <th>net</th>\n",
       "      <th>county</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>-9626792.0</td>\n",
       "      <td>3825189.75</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-9626832.0</td>\n",
       "      <td>3825073.75</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-9627101.0</td>\n",
       "      <td>3825153.50</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-9627149.0</td>\n",
       "      <td>3825322.75</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-9627159.0</td>\n",
       "      <td>3825334.75</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     easting    northing  race  net  county\n",
       "0 -9626792.0  3825189.75     1    0       0\n",
       "1 -9626832.0  3825073.75     1    0       0\n",
       "2 -9627101.0  3825153.50     1    0       0\n",
       "3 -9627149.0  3825322.75     1    0       0\n",
       "4 -9627159.0  3825334.75     1    0       0"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the dataset\n",
    "df = cudf.read_parquet('../data/total_population_dataset.parquet')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>idx</th>\n",
       "      <th>county</th>\n",
       "      <th>county_lower</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>Autauga County</td>\n",
       "      <td>autauga county</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>Baldwin County</td>\n",
       "      <td>baldwin county</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2</td>\n",
       "      <td>Barbour County</td>\n",
       "      <td>barbour county</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>Bibb County</td>\n",
       "      <td>bibb county</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>4</td>\n",
       "      <td>Blount County</td>\n",
       "      <td>blount county</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   idx          county    county_lower\n",
       "0    0  Autauga County  autauga county\n",
       "1    1  Baldwin County  baldwin county\n",
       "2    2  Barbour County  barbour county\n",
       "3    3     Bibb County     bibb county\n",
       "4    4   Blount County   blount county"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load the state to county mapping\n",
    "id2county = pickle.load(open('../id2county.pkl','rb'))\n",
    "df_counties = cudf.DataFrame(dict(idx=list(id2county.keys()), county=list(id2county.values())))\n",
    "\n",
    "# Lowercase the county names for easier merging\n",
    "df_counties['county_lower'] = df_counties.county.str.lower()\n",
    "df_counties.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>county</th>\n",
       "      <th>type</th>\n",
       "      <th>state</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Harrison</td>\n",
       "      <td>county</td>\n",
       "      <td>Missouri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Jefferson</td>\n",
       "      <td>county</td>\n",
       "      <td>Missouri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Newton</td>\n",
       "      <td>county</td>\n",
       "      <td>Missouri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Wayne</td>\n",
       "      <td>county</td>\n",
       "      <td>Missouri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Lincoln</td>\n",
       "      <td>county</td>\n",
       "      <td>Montana</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      county    type     state\n",
       "0   Harrison  county  Missouri\n",
       "1  Jefferson  county  Missouri\n",
       "2     Newton  county  Missouri\n",
       "3      Wayne  county  Missouri\n",
       "4    Lincoln  county   Montana"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Dataset downloaded from https://public.opendatasoft.com/explore/dataset/georef-united-states-of-america-county/export/?disjunctive.ste_code&disjunctive.ste_name&disjunctive.coty_code&disjunctive.coty_name\n",
    "county_state_df = cudf.read_csv('../data/us-counties1.csv', delimiter=\";\")[['Official Name County', 'Type', 'Official Name State']].dropna()\n",
    "county_state_df.columns = ['county', 'type', 'state']\n",
    "county_state_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the type to the county name\n",
    "county_state_df['county'] = county_state_df.apply(lambda row: row['county'] + ' ' + row['type'], axis=1)\n",
    "\n",
    "# Remove non-ascii characters and abbreviations to match the other id2county mapping dataset\n",
    "county_state_df['county'] = county_state_df.county.to_pandas().replace({r'[^\\x00-\\x7F]+': '', r'([A-Z][a-z]+)([A-Z]+)': r'\\1'}, regex=True)\n",
    "\n",
    "# Lowercase the county names for easier merging\n",
    "county_state_df['county_lower'] = county_state_df['county'].str.lower()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Merge the datasets and drop duplicates to get the state for each county in the total population dataset\n",
    "df_map_county_to_states = df_counties.merge(county_state_df, on='county_lower', how='left', suffixes=['', '_y']).drop_duplicates(subset=['county_lower'])[['idx', 'county', 'state' ]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Fill in the states for unavailable states manually by looking at the counties\n",
    "# Carson City, Nevada\n",
    "# District of Columbia, Washington DC\n",
    "# Remaining, Connecticut\n",
    "df_map_county_to_states.loc[df_map_county_to_states.county == 'Carson City', 'state'] = 'Nevada'\n",
    "df_map_county_to_states.loc[df_map_county_to_states.county == 'District of Columbia', 'state'] = 'Nevada'\n",
    "df_map_county_to_states.loc[df_map_county_to_states.isna().any(axis=1), 'state'] = 'Connecticut'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the mapping\n",
    "df_map_county_to_states.to_parquet('../data/county_to_state_mapping.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>idx</th>\n",
       "      <th>county</th>\n",
       "      <th>state</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>144</td>\n",
       "      <td>Lonoke County</td>\n",
       "      <td>Arkansas</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>145</td>\n",
       "      <td>Miller County</td>\n",
       "      <td>Georgia</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>146</td>\n",
       "      <td>Mississippi County</td>\n",
       "      <td>Missouri</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>147</td>\n",
       "      <td>Nevada County</td>\n",
       "      <td>California</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>148</td>\n",
       "      <td>Newton County</td>\n",
       "      <td>Texas</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2954</th>\n",
       "      <td>76</td>\n",
       "      <td>Fairbanks North Star Borough</td>\n",
       "      <td>Alaska</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2955</th>\n",
       "      <td>78</td>\n",
       "      <td>Hoonah-Angoon Census Area</td>\n",
       "      <td>Alaska</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2956</th>\n",
       "      <td>79</td>\n",
       "      <td>Juneau City and Borough</td>\n",
       "      <td>Alaska</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2958</th>\n",
       "      <td>74</td>\n",
       "      <td>Denali Borough</td>\n",
       "      <td>Alaska</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2959</th>\n",
       "      <td>77</td>\n",
       "      <td>Haines Borough</td>\n",
       "      <td>Alaska</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1955 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      idx                        county       state\n",
       "0     144                 Lonoke County    Arkansas\n",
       "1     145                 Miller County     Georgia\n",
       "3     146            Mississippi County    Missouri\n",
       "5     147                 Nevada County  California\n",
       "7     148                 Newton County       Texas\n",
       "...   ...                           ...         ...\n",
       "2954   76  Fairbanks North Star Borough      Alaska\n",
       "2955   78     Hoonah-Angoon Census Area      Alaska\n",
       "2956   79       Juneau City and Borough      Alaska\n",
       "2958   74                Denali Borough      Alaska\n",
       "2959   77                Haines Borough      Alaska\n",
       "\n",
       "[1955 rows x 3 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_map_county_to_states"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
